import torch
import torch.nn.functional as F
import os
import numpy as np

import config

os.environ['CUDA_VISIBLE_DEVICES'] = '1'


def _calc(h, t, r):
    return torch.norm(h + r - t, p=2, dim=1)


def predict(batch):
    pos_h = batch[:, 0]
    pos_r = batch[:, 1]
    pos_t = batch[:, 2]

    # target vector
    p_h = entity_emb[pos_h.cpu().numpy()]
    p_t = entity_emb[pos_t.cpu().numpy()]
    p_r = relation_emb[pos_r.cpu().numpy()]

    # score for loss
    p_score = _calc(p_h, p_t, p_r)

    return p_score


def test_head(golden_triple):
    head_batch = config.get_head_batch(golden_triple)
    value = list(predict(head_batch))
    li = np.argsort(value)
    res = 0
    sub = 0
    for pos, val in enumerate(li):
        if val == golden_triple[0]:
            res = pos + 1
            break
        if (val, golden_triple[1], golden_triple[2]) in train_set:
            sub += 1

    del head_batch
    del value
    del li

    return res, res - sub


def test_tail(golden_triple):
    tail_batch = config.get_tail_batch(golden_triple)
    value = list(predict(tail_batch))
    li = np.argsort(value)
    res = 0
    sub = 0
    for pos, val in enumerate(li):
        if val == golden_triple[2]:
            res = pos + 1
            break
        if (golden_triple[0], golden_triple[1], val) in train_set:
            sub += 1

    del tail_batch
    del value
    del li

    return res, res - sub


def test_link_prediction(test_list):
    '''
    遍历所有三元组，对于每个三元组，替换头实体为所有其他实体，再判断正确三元组所在的位置，记录下来
    替换尾实体为所有其他实体，重复同样操作
    '''
    # test_list = read_test_file()
    # test_list = read_file(train_file_name='./data/YAGO3-10-part/test2id.txt')
    test_total = len(test_list)

    l_mr = 0
    r_mr = 0
    l_hit1 = 0
    l_hit3 = 0
    l_hit10 = 0
    r_hit1 = 0
    r_hit3 = 0
    r_hit10 = 0

    l_mr_filter = 0
    r_mr_filter = 0
    l_hit1_filter = 0
    l_hit3_filter = 0
    l_hit10_filter = 0
    r_hit1_filter = 0
    r_hit3_filter = 0
    r_hit10_filter = 0

    for i, golden_triple in enumerate(test_list):
        print('test ---' + str(i) + '--- triple')
        l_pos, l_filter_pos = test_head(golden_triple)
        r_pos, r_filter_pos = test_tail(golden_triple)  # position, 1-based

        print(golden_triple, end=': ')
        print('l_pos=' + str(l_pos), end=', ')
        print('l_filter_pos=' + str(l_filter_pos), end=', ')
        print('r_pos=' + str(r_pos), end=', ')
        print('r_filter_pos=' + str(r_filter_pos), end='\n')

        l_mr += l_pos
        r_mr += r_pos

        if l_pos <= 1:
            l_hit1 += 1
        if l_pos <= 3:
            l_hit3 += 1
        if l_pos <= 10:
            l_hit10 += 1

        if r_pos <= 1:
            r_hit1 += 1
        if r_pos <= 3:
            r_hit3 += 1
        if r_pos <= 10:
            r_hit10 += 1

        ####################
        l_mr_filter += l_filter_pos
        r_mr_filter += r_filter_pos

        if l_filter_pos <= 1:
            l_hit1_filter += 1
        if l_filter_pos <= 3:
            l_hit3_filter += 1
        if l_filter_pos <= 10:
            l_hit10_filter += 1

        if r_filter_pos <= 1:
            r_hit1_filter += 1
        if r_filter_pos <= 3:
            r_hit3_filter += 1
        if r_filter_pos <= 10:
            r_hit10_filter += 1

    l_mr /= test_total
    r_mr /= test_total
    l_hit1 /= test_total
    l_hit3 /= test_total
    l_hit10 /= test_total
    r_hit1 /= test_total
    r_hit3 /= test_total
    r_hit10 /= test_total

    l_mr_filter /= test_total
    r_mr_filter /= test_total
    l_hit1_filter /= test_total
    l_hit3_filter /= test_total
    l_hit10_filter /= test_total
    r_hit1_filter /= test_total
    r_hit3_filter /= test_total
    r_hit10_filter /= test_total

    print('\t\t\tmean_rank\t\t\thit@10\t\t\thit@3\t\t\thit@1')
    print('head(raw)\t\t\t' + str(l_mr) + '\t\t\t' + str(l_hit10) + '\t\t\t' + str(l_hit3) + '\t\t\t' + str(l_hit1))
    print('tail(raw)\t\t\t' + str(r_mr) + '\t\t\t' + str(r_hit10) + '\t\t\t' + str(r_hit3) + '\t\t\t' + str(r_hit1))
    print('head(filter)\t\t\t' + str(l_mr_filter) + '\t\t\t' + str(l_hit10_filter) + '\t\t\t' + str(l_hit3_filter) + '\t\t\t' + str(l_hit1_filter))
    print('tail(filter)\t\t\t' + str(r_mr_filter) + '\t\t\t' + str(r_hit10_filter) + '\t\t\t' + str(r_hit3_filter) + '\t\t\t' + str(r_hit1_filter))


train_set = set(config.train_list)
entity_emb, relation_emb = config.load_parameter('transe_parameters100', mode='transe')
print('test link prediction starting...')
test_link_prediction(config.test_list)
print('test link prediction ending...')
